import os
import random
import numpy as np
import torch
from torch.utils.data import Dataset, Sampler


def load_data(data_path, prefix):
    assert prefix in ["train", "test"]

    data = np.load(os.path.join(data_path, f"{prefix}_data.npy"))
    token = np.load(os.path.join(data_path, f"{prefix}_token.npy"))
    label = np.load(os.path.join(data_path, f"{prefix}_label.npy"))
    size = np.load(os.path.join(data_path, f"{prefix}_size.npy"))
    return data, token, label, size


def get_window_index(data_size, window_size, stride_size, pred_size):
    input_size = window_size - stride_size
    data_index = list()
    token_index = list()
    label_index = list()
    num_windows = list()
    init_index = 0
    for i in range(len(data_size)):
        # num_window = (data_size[i] - input_size) // stride_size - 1
        num_window = (data_size[i] - input_size) // stride_size

        if num_window <= 0:
            index_start = init_index + data_size
            continue

        index_start = init_index + data_size[i] - stride_size * (num_window - 1) - window_size
        num_windows.append(num_window)
        for j in range(num_window):
            window_start = index_start + stride_size * j
            window_end = window_start + window_size
            data_index.append(list(range(window_start, window_end)))
            token_index.append(i)
            label_index.append(list(range(window_end - pred_size, window_end)))

        init_index += data_size[i]

    return data_index, token_index, label_index, num_windows


class WindowDataSet(Dataset):

    def __init__(self, data_path, select, window_size, stride_size=1, pred_size=1):
        self.data, self.token, self.label, data_size = load_data(data_path, select)
        self.data_indice, self.token_indice, self.label_indice, self.num_windows = get_window_index(data_size.tolist(),
                                                                                                    window_size,
                                                                                                    stride_size,
                                                                                                    pred_size)
        self.size = len(self.data_indice)

    def __len__(self):
        return self.size

    def __getitem__(self, index):
        n = len(self.data_indice[index])
        data = self.data[self.data_indice[index], :].copy()
        token = self.token[self.token_indice[index]].copy()
        label = self.label[self.label_indice[index]].copy()
        return data, token, label


class SeqDataSet(Dataset):

    def __init__(self, data_path, select):
        self.data, self.token, self.label, self.seq = load_data(data_path, select)
        self.size = self.seq.shape[0]
        self.indice = [0 for _ in range(self.size + 1)]
        for i in range(self.size):
            self.indice[i + 1] = self.indice[i] + self.seq[i]

    def __len__(self):
        return self.size

    def __getitem__(self, index):
        data = self.data[self.indice[index]: self.indice[index + 1], :]
        label = self.label[self.indice[index]: self.indice[index + 1], np.newaxis]
        token = self.token[index]
        return data, token, label


class BearingSampler(Sampler):

    def __init__(self, data_source):
        self.data_source = data_source
        windows = data_source.num_windows
        c = 0
        win_idx = []

        for w in windows:
            win_idx.append(range(c, c + w))
            c += w
        self.window_idx = win_idx

    def __len__(self):
        return len(self.window_idx)

    def __iter__(self):
        return iter(self.window_idx)


class PartSampler(Sampler):

    def __init__(self, data_source, start=0, end=1, shuffle=False):
        assert start >= 0 and start < 1
        assert end <= 1 and end > start
        self.data_source = data_source
        self.shuffle = shuffle
        windows = data_source.num_windows
        c = 0
        win_idx = []
        for w in windows:
            start_ = c + int(start * w)
            end_ = c + int(end * w)
            if start_ < end_:
                win_idx += list(range(start_, end_))
            c += w
        self.win_idx = win_idx

    def __iter__(self):
        if self.shuffle:
            random.shuffle(self.win_idx)
        return iter(self.win_idx)

    def __len__(self):
        return len(self.win_idx)
